"""
@author: chenzhenhua
@project: jf_fashion
@file: cifar10.py
@time: 2021/8/4 0004 0:13
@desc:
"""

import unittest

from tensorflow.python.keras.optimizers import Adam

from jf_fashion.keras.cifar10 import load_cifar10, preprocessing
from jf_fashion.keras.model import create_vgg16
from tensorflow.keras.applications.resnet50 import ResNet50


class TestCase(unittest.TestCase):

    def setUp(self) -> None:
        (x1, y1), (x2, y2) = load_cifar10()
        x_train, y_train, x_test, y_test = preprocessing(x1, y1, x2, y2)
        self.x1 = x1
        self.x2 = x2
        self.y1 = y1
        self.y2 = y2
        self.x_train = x_train
        self.y_train = y_train
        self.x_test = x_test
        self.y_test = y_test

    @unittest.skip("")
    def test_mnist_train(self):
        model = create_vgg16((32, 32, 3), 10)

        # initiate RMSprop optimizer
        opt = Adam(lr=0.0001)

        # Let's train the model using RMSprop
        model.compile(loss='categorical_crossentropy',
                      optimizer=opt,
                      metrics=['accuracy'])

        print("train____________")
        model.fit(self.x_train, self.y_train, epochs=10, batch_size=128)
        print("test_____________")
        loss, acc = model.evaluate(self.x_test, self.y_test)
        print("loss=", loss)
        print("accuracy=", acc)
        model.save('data/cifar10.h5')

    #@unittest.skip("")
    def test_model(self):
        model = ResNet50(weights='imagenet')


if __name__ == '__main__':
    unittest.main()
